import os
import shutil
import random
import re

def process_photos_with_correct_sorting(txt_file, keyword="OD202", start_count=5, end_count=5, middle_count=5, output_dir="output_photos"):
    if not os.path.exists(txt_file):
        print(f"File {txt_file} not found.")
        return

    # 创建输出文件夹及其子文件夹
    start_folder = os.path.join(output_dir, "start")
    end_folder = os.path.join(output_dir, "end")
    middle_folder = os.path.join(output_dir, "middle")
    os.makedirs(start_folder, exist_ok=True)
    os.makedirs(end_folder, exist_ok=True)
    os.makedirs(middle_folder, exist_ok=True)

    # 读取照片路径并筛选包含关键字的
    with open(txt_file, "r") as file:
        photo_paths = [line.strip() for line in file if keyword in line]

    # 按照片流分组（基于SN前缀+OD编号）
    photo_streams = {}
    for path in photo_paths:
        base_name = os.path.basename(path)
        stream_id = "_".join(base_name.split("_")[:2])  # 提取流 ID (SN+OD部分)
        photo_streams.setdefault(stream_id, []).append(path)

    # 数字化排序的辅助函数
    def sort_by_number(photo_path):
        base_name = os.path.basename(photo_path)
        match = re.search(r'_(\d+)\.jpg$', base_name)  # 提取数字序号
        return int(match.group(1)) if match else float('inf')

    # 对每个流提取指定照片并复制
    for stream_id, photos in photo_streams.items():
        photos.sort(key=sort_by_number)  # 按数值排序

        # 获取开头、结尾和中间的照片
        start_photos = photos[:start_count]
        end_photos = photos[-end_count:]
        middle_photos = random.sample(photos[start_count:-end_count], min(middle_count, len(photos[start_count:-end_count])))

        # 复制照片到对应的子文件夹
        for photo in start_photos:
            shutil.copy(photo, os.path.join(start_folder, os.path.basename(photo)))
        for photo in end_photos:
            shutil.copy(photo, os.path.join(end_folder, os.path.basename(photo)))
        for photo in middle_photos:
            shutil.copy(photo, os.path.join(middle_folder, os.path.basename(photo)))

    print(f"Photos processed and copied to {output_dir}")

# 示例用法
process_photos_with_correct_sorting(
    txt_file="gt_SN8_20241229.txt",  # 替换为你的txt文件路径
    keyword="OD202",
    start_count=5,
    end_count=5,
    middle_count=5,
    output_dir="find_bef_mid_end"
)
